#!/bin/bash

# AIQuant系统部署检查脚本
# 验证系统部署的完整性和正确性

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 检查结果统计
TOTAL_CHECKS=0
PASSED_CHECKS=0
FAILED_CHECKS=0
WARNING_CHECKS=0

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[PASS]${NC} $1"
    ((PASSED_CHECKS++))
}

log_warning() {
    echo -e "${YELLOW}[WARN]${NC} $1"
    ((WARNING_CHECKS++))
}

log_error() {
    echo -e "${RED}[FAIL]${NC} $1"
    ((FAILED_CHECKS++))
}

# 执行检查
run_check() {
    local check_name="$1"
    local check_command="$2"
    
    ((TOTAL_CHECKS++))
    log_info "检查: $check_name"
    
    if eval "$check_command" &> /dev/null; then
        log_success "$check_name"
        return 0
    else
        log_error "$check_name"
        return 1
    fi
}

# 检查文件存在性
check_file_exists() {
    local file_path="$1"
    local description="$2"
    
    ((TOTAL_CHECKS++))
    if [[ -f "$file_path" ]]; then
        log_success "$description: $file_path"
        return 0
    else
        log_error "$description: $file_path 不存在"
        return 1
    fi
}

# 检查目录存在性
check_directory_exists() {
    local dir_path="$1"
    local description="$2"
    
    ((TOTAL_CHECKS++))
    if [[ -d "$dir_path" ]]; then
        log_success "$description: $dir_path"
        return 0
    else
        log_error "$description: $dir_path 不存在"
        return 1
    fi
}

# 检查Python依赖
check_python_dependencies() {
    log_info "检查Python依赖..."
    
    local required_packages=(
        "fastapi"
        "uvicorn"
        "sqlalchemy"
        "asyncpg"
        "redis"
        "celery"
        "streamlit"
        "pandas"
        "numpy"
        "pydantic"
        "python-multipart"
        "python-jose"
        "passlib"
        "bcrypt"
        "websockets"
        "aiofiles"
        "prometheus-client"
        "psutil"
    )
    
    for package in "${required_packages[@]}"; do
        if python3 -c "import $package" 2>/dev/null; then
            log_success "Python包: $package"
            ((PASSED_CHECKS++))
        else
            log_error "Python包: $package 未安装"
            ((FAILED_CHECKS++))
        fi
        ((TOTAL_CHECKS++))
    done
}

# 检查系统服务
check_system_services() {
    log_info "检查系统服务..."
    
    local services=("docker" "nginx" "postgresql" "redis")
    
    for service in "${services[@]}"; do
        ((TOTAL_CHECKS++))
        if systemctl is-active --quiet "$service" 2>/dev/null; then
            log_success "系统服务: $service 正在运行"
            ((PASSED_CHECKS++))
        elif systemctl is-enabled --quiet "$service" 2>/dev/null; then
            log_warning "系统服务: $service 已启用但未运行"
            ((WARNING_CHECKS++))
        else
            log_error "系统服务: $service 未安装或未启用"
            ((FAILED_CHECKS++))
        fi
    done
}

# 检查Docker容器
check_docker_containers() {
    log_info "检查Docker容器..."
    
    if ! command -v docker &> /dev/null; then
        log_error "Docker未安装"
        ((FAILED_CHECKS++))
        ((TOTAL_CHECKS++))
        return 1
    fi
    
    local containers=("aiquant-api" "aiquant-web" "postgres" "redis" "nginx")
    
    for container in "${containers[@]}"; do
        ((TOTAL_CHECKS++))
        if docker ps --format "table {{.Names}}" | grep -q "$container"; then
            log_success "Docker容器: $container 正在运行"
            ((PASSED_CHECKS++))
        else
            log_warning "Docker容器: $container 未运行"
            ((WARNING_CHECKS++))
        fi
    done
}

# 检查网络连接
check_network_connectivity() {
    log_info "检查网络连接..."
    
    local endpoints=(
        "http://localhost:8000/health"
        "http://localhost:8501"
        "http://localhost:6379"
        "http://localhost:5432"
    )
    
    for endpoint in "${endpoints[@]}"; do
        ((TOTAL_CHECKS++))
        if curl -s -f "$endpoint" &> /dev/null; then
            log_success "网络连接: $endpoint"
            ((PASSED_CHECKS++))
        else
            log_warning "网络连接: $endpoint 无法访问"
            ((WARNING_CHECKS++))
        fi
    done
}

# 检查数据库连接
check_database_connection() {
    log_info "检查数据库连接..."
    
    ((TOTAL_CHECKS++))
    if python3 -c "
import asyncio
import asyncpg
import os
from dotenv import load_dotenv

async def test_connection():
    load_dotenv()
    try:
        conn = await asyncpg.connect(
            host=os.getenv('DB_HOST', 'localhost'),
            port=int(os.getenv('DB_PORT', 5432)),
            user=os.getenv('DB_USER', 'aiquant'),
            password=os.getenv('DB_PASSWORD', 'password'),
            database=os.getenv('DB_NAME', 'aiquant')
        )
        await conn.close()
        return True
    except Exception as e:
        return False

result = asyncio.run(test_connection())
exit(0 if result else 1)
" 2>/dev/null; then
        log_success "数据库连接测试"
        ((PASSED_CHECKS++))
    else
        log_error "数据库连接失败"
        ((FAILED_CHECKS++))
    fi
}

# 检查Redis连接
check_redis_connection() {
    log_info "检查Redis连接..."
    
    ((TOTAL_CHECKS++))
    if python3 -c "
import redis
import os
from dotenv import load_dotenv

load_dotenv()
try:
    r = redis.Redis(
        host=os.getenv('REDIS_HOST', 'localhost'),
        port=int(os.getenv('REDIS_PORT', 6379)),
        password=os.getenv('REDIS_PASSWORD', None),
        decode_responses=True
    )
    r.ping()
    exit(0)
except Exception as e:
    exit(1)
" 2>/dev/null; then
        log_success "Redis连接测试"
        ((PASSED_CHECKS++))
    else
        log_error "Redis连接失败"
        ((FAILED_CHECKS++))
    fi
}

# 检查API端点
check_api_endpoints() {
    log_info "检查API端点..."
    
    local api_endpoints=(
        "/health"
        "/docs"
        "/openapi.json"
        "/api/v1/health"
    )
    
    local base_url="http://localhost:8000"
    
    for endpoint in "${api_endpoints[@]}"; do
        ((TOTAL_CHECKS++))
        if curl -s -f "${base_url}${endpoint}" &> /dev/null; then
            log_success "API端点: $endpoint"
            ((PASSED_CHECKS++))
        else
            log_warning "API端点: $endpoint 无法访问"
            ((WARNING_CHECKS++))
        fi
    done
}

# 检查配置文件
check_configuration_files() {
    log_info "检查配置文件..."
    
    # 核心配置文件
    check_file_exists ".env" "环境配置文件"
    check_file_exists "requirements.txt" "Python依赖文件"
    check_file_exists "requirements-dev.txt" "开发依赖文件"
    
    # Docker配置
    check_file_exists "Dockerfile" "Docker构建文件"
    check_file_exists "docker-compose.yml" "Docker编排文件"
    check_file_exists "docker-compose.prod.yml" "生产环境Docker编排文件"
    check_file_exists ".dockerignore" "Docker忽略文件"
    
    # 应用配置
    check_file_exists "config/production.py" "生产环境配置"
    check_file_exists "config/development.py" "开发环境配置"
    check_file_exists "config/testing.py" "测试环境配置"
    check_file_exists "config/logging.conf" "日志配置"
    
    # 服务器配置
    check_file_exists "config/nginx.conf" "Nginx配置"
    check_file_exists "config/redis.conf" "Redis配置"
    
    # 监控配置
    check_file_exists "monitoring/prometheus.yml" "Prometheus配置"
    check_file_exists "monitoring/grafana-dashboard.json" "Grafana仪表板配置"
}

# 检查脚本文件
check_script_files() {
    log_info "检查脚本文件..."
    
    check_file_exists "scripts/start_system.py" "系统启动脚本"
    check_file_exists "scripts/init_db.py" "数据库初始化脚本"
    check_file_exists "scripts/health_check.py" "健康检查脚本"
    check_file_exists "scripts/setup_ssl.sh" "SSL配置脚本"
    check_file_exists "scripts/deployment_check.sh" "部署检查脚本"
    check_file_exists "Makefile" "构建脚本"
}

# 检查文档文件
check_documentation_files() {
    log_info "检查文档文件..."
    
    check_file_exists "README.md" "项目说明文档"
    check_file_exists "docs/installation_guide.md" "安装指南"
    check_file_exists "docs/user_guide.md" "用户指南"
    check_file_exists "docs/api_reference.md" "API参考文档"
    check_file_exists "docs/deployment_guide.md" "部署指南"
}

# 检查测试文件
check_test_files() {
    log_info "检查测试文件..."
    
    check_file_exists "tests/conftest.py" "测试配置文件"
    check_file_exists "tests/test_system_integration.py" "系统集成测试"
    check_file_exists "tests/test_api_endpoints.py" "API端点测试"
}

# 检查目录结构
check_directory_structure() {
    log_info "检查目录结构..."
    
    local required_dirs=(
        "aiquant"
        "aiquant/core"
        "aiquant/agents"
        "aiquant/strategies"
        "aiquant/data"
        "aiquant/api"
        "aiquant/web"
        "config"
        "scripts"
        "docs"
        "tests"
        "monitoring"
        "logs"
    )
    
    for dir in "${required_dirs[@]}"; do
        check_directory_exists "$dir" "目录"
    done
}

# 检查权限
check_permissions() {
    log_info "检查文件权限..."
    
    local executable_files=(
        "scripts/start_system.py"
        "scripts/init_db.py"
        "scripts/health_check.py"
        "scripts/setup_ssl.sh"
        "scripts/deployment_check.sh"
    )
    
    for file in "${executable_files[@]}"; do
        ((TOTAL_CHECKS++))
        if [[ -x "$file" ]]; then
            log_success "可执行权限: $file"
            ((PASSED_CHECKS++))
        else
            log_warning "缺少可执行权限: $file"
            ((WARNING_CHECKS++))
        fi
    done
}

# 检查环境变量
check_environment_variables() {
    log_info "检查环境变量..."
    
    local required_vars=(
        "DB_HOST"
        "DB_PORT"
        "DB_USER"
        "DB_PASSWORD"
        "DB_NAME"
        "REDIS_HOST"
        "REDIS_PORT"
        "SECRET_KEY"
        "ENVIRONMENT"
    )
    
    if [[ -f ".env" ]]; then
        source .env
        for var in "${required_vars[@]}"; do
            ((TOTAL_CHECKS++))
            if [[ -n "${!var}" ]]; then
                log_success "环境变量: $var"
                ((PASSED_CHECKS++))
            else
                log_error "环境变量: $var 未设置"
                ((FAILED_CHECKS++))
            fi
        done
    else
        log_error "环境配置文件 .env 不存在"
        ((FAILED_CHECKS++))
        ((TOTAL_CHECKS++))
    fi
}

# 性能检查
check_system_performance() {
    log_info "检查系统性能..."
    
    # 检查内存使用
    ((TOTAL_CHECKS++))
    local memory_usage=$(free | grep Mem | awk '{printf "%.1f", $3/$2 * 100.0}')
    if (( $(echo "$memory_usage < 80" | bc -l) )); then
        log_success "内存使用率: ${memory_usage}%"
        ((PASSED_CHECKS++))
    else
        log_warning "内存使用率过高: ${memory_usage}%"
        ((WARNING_CHECKS++))
    fi
    
    # 检查磁盘使用
    ((TOTAL_CHECKS++))
    local disk_usage=$(df / | tail -1 | awk '{print $5}' | sed 's/%//')
    if (( disk_usage < 80 )); then
        log_success "磁盘使用率: ${disk_usage}%"
        ((PASSED_CHECKS++))
    else
        log_warning "磁盘使用率过高: ${disk_usage}%"
        ((WARNING_CHECKS++))
    fi
    
    # 检查CPU负载
    ((TOTAL_CHECKS++))
    local cpu_load=$(uptime | awk -F'load average:' '{print $2}' | awk '{print $1}' | sed 's/,//')
    local cpu_cores=$(nproc)
    if (( $(echo "$cpu_load < $cpu_cores" | bc -l) )); then
        log_success "CPU负载: $cpu_load (核心数: $cpu_cores)"
        ((PASSED_CHECKS++))
    else
        log_warning "CPU负载过高: $cpu_load (核心数: $cpu_cores)"
        ((WARNING_CHECKS++))
    fi
}

# 安全检查
check_security() {
    log_info "检查安全配置..."
    
    # 检查文件权限
    ((TOTAL_CHECKS++))
    if [[ -f ".env" ]]; then
        local env_perms=$(stat -c "%a" .env)
        if [[ "$env_perms" == "600" || "$env_perms" == "644" ]]; then
            log_success "环境文件权限: $env_perms"
            ((PASSED_CHECKS++))
        else
            log_warning "环境文件权限不安全: $env_perms"
            ((WARNING_CHECKS++))
        fi
    fi
    
    # 检查密钥强度
    ((TOTAL_CHECKS++))
    if [[ -f ".env" ]]; then
        source .env
        if [[ ${#SECRET_KEY} -ge 32 ]]; then
            log_success "密钥长度: ${#SECRET_KEY} 字符"
            ((PASSED_CHECKS++))
        else
            log_error "密钥长度不足: ${#SECRET_KEY} 字符 (建议至少32字符)"
            ((FAILED_CHECKS++))
        fi
    fi
}

# 生成报告
generate_report() {
    echo ""
    echo "=================================="
    echo "       部署检查报告"
    echo "=================================="
    echo "总检查项目: $TOTAL_CHECKS"
    echo "通过: $PASSED_CHECKS"
    echo "警告: $WARNING_CHECKS"
    echo "失败: $FAILED_CHECKS"
    echo ""
    
    local success_rate=$((PASSED_CHECKS * 100 / TOTAL_CHECKS))
    echo "成功率: ${success_rate}%"
    
    if [[ $FAILED_CHECKS -eq 0 ]]; then
        if [[ $WARNING_CHECKS -eq 0 ]]; then
            log_success "系统部署完全正常！"
            return 0
        else
            log_warning "系统部署基本正常，但有一些警告需要注意"
            return 1
        fi
    else
        log_error "系统部署存在问题，需要修复失败的检查项"
        return 2
    fi
}

# 主函数
main() {
    echo "AIQuant系统部署检查"
    echo "===================="
    echo ""
    
    # 执行所有检查
    check_directory_structure
    check_configuration_files
    check_script_files
    check_documentation_files
    check_test_files
    check_permissions
    check_environment_variables
    check_python_dependencies
    check_system_services
    check_docker_containers
    check_database_connection
    check_redis_connection
    check_network_connectivity
    check_api_endpoints
    check_system_performance
    check_security
    
    # 生成报告
    generate_report
}

# 运行主函数
main "$@"